package org.finesys.ai.service.impl;

import java.util.List;

import org.finesys.ai.core.chat.ChatCompletionRequest;
import org.finesys.ai.core.chat.ChatCompletionResponse;
import org.finesys.ai.core.chat.Role;
import org.finesys.ai.core.image.GenerateImagesResponse;
import org.finesys.ai.dto.ChatImageReq;
import org.finesys.ai.dto.ChatReq;
import org.finesys.ai.entity.AigcApp;
import org.finesys.ai.entity.AigcConversation;
import org.finesys.ai.entity.AigcMessages;
import org.finesys.ai.entity.AigcModel;
import org.finesys.ai.service.AigcAppService;
import org.finesys.ai.service.AigcChatService;
import org.finesys.ai.service.AigcConversationService;
import org.finesys.ai.service.AigcMessagesService;
import org.finesys.ai.service.AigcModelService;
import org.finesys.chat.core.base.chat.ChatRequest;
import org.finesys.chat.core.base.factory.ChatModelFactory;
import org.finesys.common.core.util.ServletUtil;
import org.finesys.common.core.util.ValidationUtil;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.BeanUtils;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

import com.baomidou.mybatisplus.core.toolkit.Wrappers;

import lombok.RequiredArgsConstructor;
import reactor.core.publisher.Flux;

@Service
@RequiredArgsConstructor
public class AigcChatServiceImpl implements AigcChatService {

    private final AigcModelService aigcModelService;
    private final AigcConversationService aigcConversationService;
    private final AigcAppService aigcAppService;
    private final AigcMessagesService aigcMessagesService;
    private final ChatModelFactory chatModelFactory;

    @Override
    public Flux<ChatCompletionResponse> chatFluxCompletion(ChatReq chatReq) {
        //模型信息
        String modelId = chatReq.getModelId();
        if (!StringUtils.hasText(modelId)) {
            throw new IllegalArgumentException("模型ID不能为空");
        }
        AigcModel aigcModel = aigcModelService.getById(modelId);
        ValidationUtil.ensureNotNull(aigcModel, "模型不存在");
        //构建ChatRequest对象
        ChatRequest chatRequest = getChatRequest(chatReq, aigcModel, modelId);
        //调用模型会话
        return chatModelFactory.getChatClient(modelId)
                .chatCompletion(chatRequest);
    }


    /**
     * 根据传入的ChatReq对象构建ChatCompletionRequest对象
     *
     * @param chatReq ChatReq对象，包含构建ChatCompletionRequest所需的信息
     * @return ChatCompletionRequest对象，若构建失败则返回null
     */
    private ChatCompletionRequest buildChatCompletionRequest(ChatReq chatReq) {
        //构建ChatCompletionRequest对象
        ChatCompletionRequest.Builder builder = ChatCompletionRequest.builder();
        //应用信息
        String appId = chatReq.getAppId();
        if (StringUtils.hasText(appId)) {
            AigcApp aigcApp = aigcAppService.getById(appId);
            if (aigcApp != null) {
                //模型信息
                String modelId = aigcApp.getModelId();
                if (StringUtils.hasText(modelId)) {
                    AigcModel aigcModel = aigcModelService.getById(modelId);
                    if (aigcModel != null) {
                        builder.model(aigcModel.getModelName());
                        chatReq.setModelId(modelId);
                    }
                } else {
                    throw new IllegalArgumentException("模型ID不能为空");
                }
                builder.addUserMessage(aigcApp.getPrompt());
                chatReq.setPromptText(aigcApp.getPrompt());
                chatReq.setKnowledgeIds(aigcApp.getKnowledgeIds());
            }
        }
        //用户角色
        chatReq.setRole(Role.USER.name());
        //对话信息
        String conversationId = chatReq.getConversationId();
        if (StringUtils.hasText(conversationId)) {
            AigcConversation aigcConversation = aigcConversationService.getById(conversationId);
            if (aigcConversation != null) {
                builder.addAssistantMessage(aigcConversation.getTitle());
            }
        }
        //保存会话信息
        saveMessages(chatReq, 0, 0);
        return builder.build();
    }

    /**
     * 保存消息
     *
     * @param chatReq      聊天请求对象
     * @param promptTokens 提示词数量
     * @param tokens       生成文本数量
     */
    private void saveMessages(ChatReq chatReq, Integer promptTokens, Integer tokens) {
        if (chatReq.getConversationId() != null) {
            AigcMessages aigcMessages = new AigcMessages();
            BeanUtils.copyProperties(chatReq, aigcMessages);
            aigcMessages.setIp(ServletUtil.getClientIP());
            aigcMessages.setPromptTokens(promptTokens);
            aigcMessages.setTokens(tokens);
            aigcMessagesService.save(aigcMessages);
        }
    }

    @Override
    public Flux<GenerateImagesResponse> chatImage(ChatImageReq chatImageReq) {
        //模型信息
        String modelId = chatImageReq.getModelId();
        if (!StringUtils.hasText(modelId)) {
            throw new IllegalArgumentException("模型ID不能为空");
        }
        AigcModel aigcModel = aigcModelService.getById(modelId);
        ValidationUtil.ensureNotNull(aigcModel, "模型不存在");
        //构建ChatRequest对象
        ChatRequest chatRequest = getChatRequest(chatImageReq, aigcModel, modelId);
        //调用模型生成图像
        return chatModelFactory.getChatClient(modelId).imageGeneration(chatRequest);
    }

    /**
     * 根据传入的参数构建一个ChatRequest对象
     *
     * @param chatImageReq 包含聊天图像请求参数的对象
     * @param aigcModel    包含AIGC模型参数的对象
     * @param modelId      模型ID
     * @return 构造好的ChatRequest对象
     * @throws NullPointerException 如果传入的参数为null
     */
    private static @NotNull ChatRequest getChatRequest(ChatImageReq chatImageReq, AigcModel aigcModel, String modelId) {
        ChatRequest chatRequest = new ChatRequest();
        chatRequest.setImageSize(chatImageReq.getSize());
        chatRequest.setImageStyle(chatImageReq.getStyle());
        chatRequest.setImageQuality(chatImageReq.getQuality());
        chatRequest.setModel(aigcModel.getModelName());
        chatRequest.setModelId(modelId);
        chatRequest.setChatId(chatImageReq.getChatId());
        chatRequest.setConversationId(chatImageReq.getConversationId());
        chatRequest.setTopP(aigcModel.getTopP());
        chatRequest.setTemperature(aigcModel.getTemperature());
        chatRequest.setUserMessage(chatImageReq.getPromptText());
        return chatRequest;
    }

    /**
     * 根据传入的参数构建一个ChatRequest对象
     *
     * @param chatReq   包含请求参数的对象
     * @param aigcModel 包含AIGC模型参数的对象
     * @param modelId   模型ID
     * @return 构造好的ChatRequest对象
     * @throws NullPointerException 如果传入的参数为null
     */
    private static @NotNull ChatRequest getChatRequest(ChatReq chatReq, AigcModel aigcModel, String modelId) {
        ChatRequest chatRequest = new ChatRequest();
        chatRequest.setChatId(chatReq.getChatId());
        chatRequest.setUser(chatReq.getUserId());
        chatRequest.setModelId(modelId);
        chatRequest.setModel(aigcModel.getModelName());
        chatRequest.setConversationId(chatReq.getConversationId());
        chatRequest.setKnowledgeId(chatReq.getKnowledgeId());
        chatRequest.setKnowledgeIds(chatReq.getKnowledgeIds());
        chatRequest.setDocId(chatReq.getDocId());
        chatRequest.setTopP(aigcModel.getTopP());
        chatRequest.setUserMessage(chatReq.getPromptText());
        chatRequest.setTemperature(aigcModel.getTemperature());
        chatRequest.setMaxTokens(chatReq.getMaxTokens());
        chatRequest.setStream(chatReq.getStream());
        if (chatReq.getStream() == null) {
            chatRequest.setStream(true);
        }
        String message = chatReq.getMessage();
        if (!StringUtils.hasText(chatReq.getMessage())) {
            message = chatReq.getPromptText();
        }
        if (StringUtils.hasText(chatReq.getRole()) && StringUtils.hasText(message)) {
            String role = chatReq.getRole();
            switch (role) {
                case "user":
                    chatRequest.setUserMessage(message);
                    break;
                case "assistant":
                    chatRequest.setAssistantMessage(message);
                    break;
                default:
                    chatRequest.setSystemMessage(message);
            }
        }
        return chatRequest;
    }

    /**
     * 获取模型信息
     */
    @Override
    public List<AigcModel> models(String provider) {
        List<AigcModel> aigcModelList = aigcModelService.list(Wrappers.<AigcModel>lambdaQuery().eq(AigcModel::getProvider, provider));
        aigcModelList.forEach(model -> {
            model.setApiKey("");
            model.setSecretKey("");
        });
        return aigcModelList;
    }
}
